#include "jerasure.h"

JER_Gen_T::JER_Gen_T(){
	M = NULL;
	rs_r6 = false;
	PDrive = false;
}

JER_Gen_T::~JER_Gen_T(){
	Delete_Schedules();
}

/********* Schedule functions ************/

int JER_Gen_T::Create_Encode_Schedule(int smart){

	JER_Schedule * sched;
	map <string, JER_Schedule * >::iterator it_scheds;
	string str_drive_states; //string of 0's and 1's
	int i;

	//only works for bitmatrix
	if(M == NULL || M->W != 1){
		return -1;	
	}

	//the encode schedule is represented in the map as k 0's and m 1's
	str_drive_states = "";
	for(i=0;i<K;i++){
		str_drive_states += "0";
	}
	for(i=K;i<N;i++){
		str_drive_states += "1";
	}

	//see if we already have an encode schedule in our map
	it_scheds = Schedules.find(str_drive_states);

	if(it_scheds != Schedules.end()){
		//we already have an encode schedule, but we will delete it because it could be different than the one just requested (smart vs. dumb)
		delete(it_scheds->second);
		Schedules.erase(it_scheds);
	}

	//create the schedule
	if(smart){
		sched = M->To_Schedule_CSHR(WPD); 
	}else{
		sched = M->To_Schedule_Dumb(WPD); 
	}

	if(sched != NULL){
		Schedules.insert(make_pair(str_drive_states,sched));
		return 0;
	}else{
		return -1;
	}
}

/*
	For RAID6 generator matrices, we can store all possible schedules for decoding. This creates all decoding schedules, and the encoding schedule.
*/
int JER_Gen_T::Create_R6_Schedules(int smart){

	int m;
	int k;
	int i,j;
	vector <int> erased;
	map <string, JER_Schedule * >::iterator it_scheds; //cache containing encode and decode schedules, key is a string of 0's and 1's representing which drives failed
	JER_Schedule * sched;
	string str_sched_key;
	int e1,e2;
	JER_Schedule * sched_enc;

	m = N-K;
	k = K; 

	//only works for m = 2
	if(m != 2) return -1;

	if(M == NULL){
		return -1;
	}

	if(Schedules.size() != 0){
		//we already have schedules in the map, delete them
		
		//delete all schedules in the map
		Delete_Schedules();			

	}

	//erased starts w/ all drives up
	erased.resize(N,0);
	str_sched_key = "";
	str_sched_key.resize(N,0);

	//Create all combinations of losing 1 or 2 drives
	for(i=0;i<N;i++){
		for(j=i;j<N;j++){
			erased[i] = 1;
			str_sched_key[i] = '1';
			erased[j] = 1;
			str_sched_key[j] = '1';

			sched = Create_Single_Decode_Schedule(erased, smart);
			if(sched == NULL){
				return -1;
			}
			//add the schedule to the map
			Schedules.insert(make_pair(str_sched_key,sched));
			
			erased[i] = 0;
			str_sched_key[i] = '0';
			erased[j] = 0;
			str_sched_key[j] = '0';
		}
	}

	return 0;
}

JER_Schedule * JER_Gen_T::Create_Single_Decode_Schedule(vector <int> & erased, int smart){

	int i, j, x, drive, y, index, z;
	int r,c;
	JER_Matrix *decoding_matrix, *real_decoding_matrix;
	JER_Matrix *decoding_matrix_inv;
	int *row_ids;
	int *ind_to_row;
	int ddf, cdf;
	JER_Schedule *schedule;
	int *b1, *b2;
	int xor_result;
	int b1_row;
	int b2_row;

	int k, m;

	if(M == NULL){
		return NULL;
	}

	if(M->W != 1){
		return NULL;
	}
	
	k = K;
	m = N-K;

	if(!k || !m)return NULL;

	/* First, figure out the number of data drives that have failed, and the number of coding drives that have failed: ddf and cdf */
	ddf = 0;
	cdf = 0;
	for (i = 0; i < erased.size(); i++) {
		if(erased[i] == 1){
			if(i<K){
				ddf++;
			}else{
				cdf++;
			}
		}
	}

	row_ids = talloc(int, k+m);
	ind_to_row = talloc(int, k+m);
	//setup id's for scheduled decoding:
	//See set_up_ptrs_for_scheduled_decoding for how these are set 
	j = k;
	x = k;
	//loop through data drives (id=0 to k-1)
	for (i = 0; i < k; i++) {
		if (erased[i] == 0) {
			//data drive i is not erased
			row_ids[i] = i;
			ind_to_row[i] = i;
		} else {
			//data drive i is erased
			while (erased[j]) j++;
			//j is now the lowest unused non-failed coding drive.
			row_ids[i] = j;
			ind_to_row[j] = i;
			j++;
			row_ids[x] = i;
			ind_to_row[i] = x;
			x++;
		}
	}
	//loop through the coding drives (id = k to k+m-1)
	for (i = k; i < k+m; i++) {
		if (erased[i]) {
			row_ids[x] = i;
			ind_to_row[i] = x;
			x++;
		}
	}

	/* Now, we're going to create one decoding matrix which is going to 
	   decode everything with one call.  The hope is that the scheduler
	   will do a good job.    This matrix has WPD*e rows, where e is the
	   number of erasures (ddf+cdf) */
	real_decoding_matrix = new JER_Matrix((cdf+ddf)*WPD,k*WPD,1);

	/* First, if any data drives have failed, then initialize the first
	   ddf*WPD rows of the decoding matrix from the standard decoding
	   matrix inversion */
	if (ddf > 0) {

		decoding_matrix = new JER_Matrix(k*WPD,k*WPD,1);

		//loop through all k super rows
		//i is the index of the super-row
		for (i = 0; i < k; i++) {
			if (row_ids[i] == i) {
				for(r=i*WPD;r<i*WPD+WPD;r++){
					for(c=0;c<k*WPD;c++){
						//setup an identity matrix in super-cell i,i
						if(r == c){
							decoding_matrix->Set(r,c,1);
						}
					}
				}
			} else {
				//the ith super-row should be the coding matrix's (row_ids[i]-k) super-row 
				//copy all WPD rows making up the super-row
				decoding_matrix->Copy_Panel(M,(row_ids[i]-k)*WPD,0,i*WPD,0,WPD,K*WPD);
			}
		}

		decoding_matrix_inv = Inverse(decoding_matrix);
		if(decoding_matrix_inv == NULL){
			delete real_decoding_matrix;
			delete decoding_matrix;
			return NULL;
		}
		//set the first ddf super-rows of real_decoding matrix equal to
		//the decoding_matrix's row_ids[k+i] super-row
		for (i = 0; i < ddf; i++) {
			//WPD rows make up a super-row
			real_decoding_matrix->Copy_Panel(decoding_matrix_inv,row_ids[k+i]*WPD,0,i*WPD,0,WPD,K*WPD);
		}
		delete decoding_matrix;
		delete decoding_matrix_inv;
	}

	/* Next, here comes the hard part.  For each coding node that needs
	   to be decoded, you start by putting its rows of the distribution
	   matrix into the decoding matrix.  If there were no failed data
	   nodes, then you're done.  However, if there have been failed
	   data nodes, then you need to modify the columns that correspond
	   to the data nodes.  You do that by first zeroing them.  Then
	   whereever there is a one in the distribution matrix, you XOR
	   in the corresponding row from the failed data node's entry in
	   the decoding matrix.  The whole process kind of makes my head
	   spin, but it works.
	 */
	for (x = 0; x < cdf; x++) {
		drive = row_ids[x+ddf+k]-k;

		//set real_decoding_matrix's (ddf+x)'th super-row equal to bitmatrix's drive'th super-row
		real_decoding_matrix->Copy_Panel(M,WPD*drive,0,WPD*(ddf+x),0,WPD,K*WPD);

		//loop through all super-columns
		for (i = 0; i < k; i++) {
			if (row_ids[i] != i) {

				//set real_decoding_matrix's super-cell (ddf+x, i) = 0
				for(r=0;r<WPD;r++){
					for(c=0;c<WPD;c++){
						real_decoding_matrix->Set((ddf+x)*WPD+r,i*WPD+c,0);
					}
				}
			}
		}

		// There's the yucky part 
		//index is start of super-row drive
		//index = drive*k*WPD*WPD;
		//loop through super-columns
		for (i = 0; i < k; i++) {
			if (row_ids[i] != i) {
				//b1 is the start of super-row ind_to_row[i]-k
				//b1 = real_decoding_matrix+(ind_to_row[i]-k)*k*WPD*WPD;
				b1_row = (ind_to_row[i]-k)*WPD;
				for (j = 0; j < WPD; j++) {
					b2_row = (ddf+x)*WPD+j;
					for (y = 0; y < WPD; y++) {
						if(M->Get(drive*WPD+j,i*WPD+y)){
							//loop through all the columns
							for (z = 0; z < k*WPD; z++) {
								xor_result = real_decoding_matrix->Get(b2_row,z) ^ real_decoding_matrix->Get(b1_row+y,z);
								real_decoding_matrix->Set(b2_row,z,xor_result);
							}
						}
					}
				}
			}
		}
	}

	schedule = new JER_Schedule;

	if (smart) {
		schedule = real_decoding_matrix->To_Schedule_CSHR(WPD);
	} else {
		schedule = real_decoding_matrix->To_Schedule_Dumb(WPD);
	}
	free(row_ids);
	free(ind_to_row);
	delete real_decoding_matrix;

	return schedule;
}

//Deletes all schedules for this generator
//This is called before a generator is overwritten.
void JER_Gen_T::Delete_Schedules(){

	map <string, JER_Schedule * >::iterator it_scheds;

	for(it_scheds = Schedules.begin();it_scheds != Schedules.end(); it_scheds++){
		delete it_scheds->second;
	}

	Schedules.clear();
}

/*
Returns: 1 if the matrix is MDS
0 if it is not MDS
-1 on error		
 */
int JER_Gen_T::Am_I_MDS(){

	int i;
	int m;
	vector <int> data_survivors;
	vector <int> combo;
	vector <int> prev_m_indices; //contains K elements. prev_m_indices[i] = the drive id in the ith row of m_deleted_rows
	vector <int> super_rows; //contains indices of all super-rows
	JER_Matrix *m_deleted_rows; //matrix of survivors (with data rows deleted) 
	JER_Matrix m_opt; //optimized matrix (m_deleted_rows after we delete columns corresponding to data survivors)
	JER_Matrix *m_inv; //inverse of m_tmp
	int row_id;
	int num_rows_changed; //how many rows we must change from the last combination
	int combo_index;
	int dest_row;
	int num_coding_survivors;
	int prev_num_used;

	m = N-K;

	m_deleted_rows = new JER_Matrix(K*WPD,K*WPD,M->W);
	if(!m_deleted_rows){
		return -1;
	}

	//do all combinations of deleting m super-rows 
	//(n choose k, order doesn't matter)

	//fill arr_super_rows with indices of super-rows
	for(i=0;i<N;i++){
		super_rows.push_back(i);
	}
	//create first combination
	for(i=0;i<K;i++){
		data_survivors.push_back(i);
		combo.push_back(i);
		prev_m_indices.push_back(-1);
	}

	//systematic codes have the optimization where we can delete rows and cols corresponding to saved data
	if(true){

		//skip the first combo because all the survivors are data drives (we have an identity matrix)
		num_rows_changed =next_combination(super_rows,combo);

		while(num_rows_changed != 0){
			prev_num_used = K-data_survivors.size();	
			if(data_survivors.size() > K-num_rows_changed){
				data_survivors.resize(K-num_rows_changed);
			}
			dest_row = prev_num_used-num_rows_changed;//the row in m_deleted_rows we are copying into 
			if(dest_row < 0){
				dest_row = 0;
			}

			//only change the last 'K-num_rows_changed' rows
			for(combo_index = K-num_rows_changed; combo_index < K;combo_index++){

				row_id = combo[combo_index];

				if(row_id < K){
					//a data row survived
					data_survivors.push_back(row_id);
				}else{
					//coding row survived
					if(prev_m_indices[dest_row] != row_id){
						//we need to modify the matrix row because it changed from last time

						//take a (super)row from the generator matrix
						/*
						   source row = row_id * WPD
						   source col = 0;
						   dest row = dest_row * WPD
						   dest col = 0
						   number of rows = WPD
						   number of cols = WPD * K
						 */
						prev_m_indices[dest_row] = row_id; 

						row_id-=K;

						m_deleted_rows->Copy_Panel(M,row_id*WPD,0,dest_row*WPD,0,WPD,WPD*K); 

					}
					dest_row++;
				}
			}

			//if we have data survivors,
			//create a copy of m_deleted_rows and delete columns
			//this creates an optimized matrix
			if(data_survivors.size()!=0){
				num_coding_survivors = K-data_survivors.size();

				//the size of m_deleted_rows is always K by K, but we may not want all K rows
				m_opt.Create_Empty(num_coding_survivors*WPD,K*WPD,M->W);
				m_opt.Copy_Panel(m_deleted_rows,0,0,0,0,num_coding_survivors*WPD,K*WPD); 

				//delete columns corresponding to data survivors
				for(i=0;i<(K-num_coding_survivors);i++){
					m_opt.Delete_Cols(data_survivors[i]*WPD-i*WPD,WPD);
				}
				m_inv = Inverse(&m_opt);

			}else{
				//no columns to delete
				m_inv = Inverse(m_deleted_rows);
			}

			if(m_inv == NULL){
				delete m_deleted_rows;
				return 0;
			}else{
				delete m_inv;
			}

			//get the next combination
			num_rows_changed =next_combination(super_rows,combo);

		}

	}else{
		//non-systematic
		//the code is the same as the systematic code, but the optimization
		//of not adding surviving data rows is not used
		num_rows_changed = K;
		while(num_rows_changed != 0){
			dest_row = K-num_rows_changed;//the row in m_deleted_rows we are copying into 
			if(dest_row < 0){
				dest_row = 0;
			}

			//only change the last 'K-num_rows_changed' rows
			for(combo_index = K-num_rows_changed; combo_index < K;combo_index++){

				row_id = combo[combo_index];

				if(prev_m_indices[dest_row] != row_id){
					//we need to modify the matrix row because it changed from last time

					//take a (super)row from the generator matrix
					/*
					   source row = row_id * WPD
					   source col = 0;
					   dest row = dest_row * WPD
					   dest col = 0
					   number of rows = WPD
					   number of cols = WPD * K
					 */
					prev_m_indices[dest_row] = row_id; 

					m_deleted_rows->Copy_Panel(M,row_id*WPD,0,dest_row*WPD,0,WPD,WPD*K); 

				}
				dest_row++;
			}

			m_inv = Inverse(m_deleted_rows);

			if(m_inv == NULL){
				delete m_deleted_rows;
				return 0;
			}else{
				delete m_inv;
			}

			//get the next combination
			num_rows_changed =next_combination(super_rows,combo);

		}

	}

	delete m_deleted_rows;
	return 1;
}

/*
Description:
Non-recursive function to get the next combination of N choose K rows that are not deleted

Return value:
int - the number of rows changed from the last combination
The combination vector's values always change from back to front.

Example:
for n vector: 0,1,2,3,4,5,6

if r vector: 0,1,2,3
After calling this function:
r vector: 0,1,2,4
function returns 1

if r vector: 0,1,2,6
After calling this function:
r vector: 0,1,3,4
function returns 2

 */
int JER_Gen_T::next_combination(vector <int> & n, vector <int> & r){

	bool boolmarked=false;

	int c,i;
	vector <int>::iterator r_begin;
	vector <int>::iterator r_end;
	vector <int>::iterator n_begin;
	vector <int>::iterator n_end;

	vector <int>::iterator r_marked,n_marked;
	vector <int>::iterator n_it1, n_it2, n_it3, n_it4;
	vector <int>::iterator tmp_r_end;
	vector <int>::iterator r_it1,r_it2;

	r_begin = r.begin();
	r_end = r.end();
	n_begin = n.begin();
	n_end = n.end();

	n_it1=n_end;
	--n_it1;
	tmp_r_end=r_end;
	--tmp_r_end;

	for(r_it1=tmp_r_end; r_it1!=r_begin || r_it1==r_begin; --r_it1,--n_it1){

		if(*r_it1==*n_it1 ){

			//to ensure not at the start of r sequence
			if(r_it1!=r_begin){

				boolmarked=true;
				r_marked=(--r_it1);
				++r_it1;//add it back again 

				continue;

			}else{ 
				// it means it is at the start the sequence, so return -1
				return 0;
			}
		}else{

			//marked code

			if(boolmarked==true){

				//for loop to find which marked is in the first sequence
				n_marked;//mark in first sequence

				for (n_it2=n_begin;n_it2!=n_end;++n_it2){

					if(*r_marked==*n_it2){
						n_marked=n_it2;
						break;
					}
				}

				n_it3=++n_marked;
				c=0;
				for(r_it2=r_marked;r_it2!=r_end;++r_it2,++n_it3){
					*r_it2=*n_it3;
					c++;
				}
				return c; //multiple (c) numbers of the combination changed
			}

			for(n_it4=n_begin; n_it4!=n_end; ++n_it4){

				if(*r_it1==*n_it4){
					*r_it1=*(++n_it4);
					return 1; //only the last number of the combination changed           
				}
			}
		}
	}  

	return 0;//will never reach here    
}

JER_Gen_T *JER_Gen_T::Genmatrix_To_Genbitmatrix()
{
	JER_Gen_T *result;

	result = new JER_Gen_T();
	if( result == NULL )
		return NULL;

	result->M = M->Matrix_To_Bitmatrix();
	if( result->M == NULL )
	{
		delete result;
		return NULL;
	}
	result->N = N;
	result->K = K;
	result->WPD = M->W;
	result->Systematic = Systematic;
	result->PDrive = PDrive;
	result->Schedules = Schedules;

	return result;
}

JER_Gen_T *JER_Gen_T::Genbitmatrix_To_Genmatrix()
{
	JER_Gen_T *result;

	result = new JER_Gen_T();
	if( result == NULL )
		return NULL;

	result->M = M->Bitmatrix_To_Matrix( WPD );
	if( result->M == NULL )
	{
		delete result;
		return NULL;
	}
	result->N = N;
	result->K = K;
	result->WPD = 1;
	result->Systematic = Systematic;
	result->PDrive = PDrive;
	result->Schedules = Schedules;

	return result;
}

int JER_Gen_T::Genmatrix_To_Genbitmatrix( JER_Gen_T &bgen )
{
	if( M->Matrix_To_Bitmatrix( *(bgen.M) ) )
		return -1;

	bgen.N = N;
	bgen.K = K;
	bgen.WPD = M->W;

	bgen.Systematic = Systematic;
	bgen.PDrive = PDrive;
	bgen.Schedules = Schedules;

	return 0;
}

int JER_Gen_T::Genbitmatrix_To_Genmatrix( JER_Gen_T &gen )
{
	if( M->Bitmatrix_To_Matrix( *(gen.M), WPD ) )
		return -1;

	gen.N = N;
	gen.K = K;
	gen.WPD = 1;

	gen.Systematic = Systematic;
	gen.PDrive = PDrive;
	gen.Schedules = Schedules;

	return 0;
}
